package org.uranus.webauth.fido.uranus.service.impl;

import java.io.IOException;
import java.io.InputStream;
import java.net.MalformedURLException;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;
import java.time.Clock;
import java.util.*;
import java.util.concurrent.ExecutionException;

import javax.annotation.PostConstruct;

import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.uranus.webauth.fido.uranus.cache.CacheManager;
import org.uranus.webauth.fido.uranus.model.RegistrationRequestEntity;
import org.uranus.webauth.fido.uranus.util.UranusUtil;

import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.JsonSerializer;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializerProvider;
import com.fasterxml.jackson.databind.annotation.JsonSerialize;
import com.google.common.collect.Lists;
import com.google.common.io.Closeables;
import com.upokecenter.cbor.CBORObject;
import com.yubico.internal.util.CertificateParser;
import com.yubico.internal.util.ExceptionUtil;
import com.yubico.internal.util.JacksonCodecs;
import com.yubico.webauthn.*;
import com.yubico.webauthn.attestation.*;
import com.yubico.webauthn.attestation.resolver.CompositeAttestationResolver;
import com.yubico.webauthn.attestation.resolver.CompositeTrustResolver;
import com.yubico.webauthn.attestation.resolver.SimpleAttestationResolver;
import com.yubico.webauthn.attestation.resolver.SimpleTrustResolverWithEquality;
import com.yubico.webauthn.data.*;
import com.yubico.webauthn.exception.RegistrationFailedException;
import com.yubico.webauthn.extension.appid.InvalidAppIdException;

import demo.webauthn.Config;
import demo.webauthn.InMemoryRegistrationStorage;
import demo.webauthn.RegistrationStorage;
import demo.webauthn.WebAuthnServer;
import demo.webauthn.data.*;
import lombok.AllArgsConstructor;
import lombok.NonNull;
import lombok.Value;

/**
 * 修改后的流程
 * 
 * @author lwl
 * @date 2020/5/25 17:17
 */
@Component
public class CustomWebAuthServer {

	private static final Logger logger = LoggerFactory.getLogger(CustomWebAuthServer.class);

	private static final String PREVIEW_METADATA_PATH = "/preview-metadata.json";

	private TrustResolver trustResolver = null;

	private MetadataService metadataService = null;

	private final Clock clock = Clock.systemDefaultZone();

	private final ObjectMapper jsonMapper = JacksonCodecs.json();

	@Autowired
	private CacheManager cacheManager;

	@PostConstruct
	void init() throws InvalidAppIdException, CertificateException {
		trustResolver = new CompositeTrustResolver(Arrays.asList(
				StandardMetadataService.createDefaultTrustResolver(),
				createExtraTrustResolver()));

		metadataService = new StandardMetadataService(
				new CompositeAttestationResolver(Arrays.asList(
						StandardMetadataService.createDefaultAttestationResolver(trustResolver),
						createExtraMetadataResolver(trustResolver))));
	}

	private RelyingParty getRelyingParty(RegistrationStorage userStorage) {
		RelyingParty rp = RelyingParty.builder()
				.identity(Config.getRpIdentity())
				.credentialRepository(userStorage)
				.origins(Config.getOrigins())
				.attestationConveyancePreference(Optional.of(AttestationConveyancePreference.NONE))
				.metadataService(Optional.of(metadataService))
				.allowOriginPort(false)
				.allowOriginSubdomain(false)
				.allowUnrequestedExtensions(true)
				.allowUntrustedAttestation(true)
				.validateSignatureCounter(false)
				.appId(Config.getAppId())
				.build();
		return rp;
	}

	private RegistrationStorage getUserStorage() {
		RegistrationStorage userStorage = new InMemoryRegistrationStorage();
		return userStorage;
	}

	private RegistrationStorage getAssertUserStorage(String username, String responses,
			ByteArray userHandle) {
		RegistrationStorage userStorage = new InMemoryRegistrationStorage();
		RegistrationResponse response = null;
		try{
			response = jsonMapper.readValue(responses, RegistrationResponse.class);
		} catch (IOException e){
			logger.error("JSON error in finishRegistration; responseJson: {}", responses, e);
		}
		CredentialRegistration reg = CredentialRegistration.builder()
				.userIdentity(
						UserIdentity.builder().name(username).displayName(username).id(userHandle)
								.build())
				.credentialNickname(null)
				.registrationTime(clock.instant())
				.credential(RegisteredCredential.builder()
						.credentialId(response.getCredential().getResponse().getAttestation()
								.getAuthenticatorData()
								.getAttestedCredentialData().get().getCredentialId())
						.userHandle(userHandle)
						.publicKeyCose(response.getCredential().getResponse().getAttestation()
								.getAuthenticatorData()
								.getAttestedCredentialData().get().getCredentialPublicKey())
						.signatureCount(
								response.getCredential().getResponse().getParsedAuthenticatorData()
										.getSignatureCounter())
						.build())
				.signatureCount(0)
				.attestationMetadata(null)
				.build();
		userStorage.addRegistrationByUsername(username, reg);
		return userStorage;
	}

	private static MetadataObject readPreviewMetadata() {
		InputStream is = WebAuthnServer.class.getResourceAsStream(PREVIEW_METADATA_PATH);
		try{
			return JacksonCodecs.json().readValue(is, MetadataObject.class);
		} catch (IOException e){
			throw ExceptionUtil
					.wrapAndLog(logger, "Failed to read metadata from " + PREVIEW_METADATA_PATH, e);
		} finally{
			Closeables.closeQuietly(is);
		}
	}

	/**
	 * Create a {@link TrustResolver} that accepts attestation certificates that are
	 * directly recognised as trust anchors.
	 */
	private static TrustResolver createExtraTrustResolver() {
		try{
			MetadataObject metadata = readPreviewMetadata();
			return new SimpleTrustResolverWithEquality(metadata.getParsedTrustedCertificates());
		} catch (CertificateException e){
			throw ExceptionUtil.wrapAndLog(logger, "Failed to read trusted certificate(s)", e);
		}
	}

	/**
	 * Create a {@link AttestationResolver} with additional metadata for unreleased
	 * YubiKey Preview devices.
	 */
	private static AttestationResolver createExtraMetadataResolver(TrustResolver trustResolver) {
		try{
			MetadataObject metadata = readPreviewMetadata();
			return new SimpleAttestationResolver(Collections.singleton(metadata), trustResolver);
		} catch (CertificateException e){
			throw ExceptionUtil.wrapAndLog(logger, "Failed to read trusted certificate(s)", e);
		}
	}

	private RegistrationRequest getRegistrationRequest(
			RegistrationRequestEntity registrationRequestEntity) {

		final UserIdentity registrationUserId = UserIdentity.builder()
				.name(registrationRequestEntity.getUsername())
				.displayName(registrationRequestEntity.getDisplayName())
				.id(new ByteArray(
						Base64.getDecoder().decode(registrationRequestEntity.getUserIdentityId())))
				.build();

		RegistrationRequest request = new RegistrationRequest(
				registrationRequestEntity.getUsername(),
				Optional.of(registrationRequestEntity.getCredentialNickname()),
				new ByteArray(Base64.getDecoder().decode(registrationRequestEntity.getChallenge())),
				this.getRelyingParty(getUserStorage()).startRegistration(
						StartRegistrationOptions.builder()
								.user(registrationUserId)
								.authenticatorSelection(AuthenticatorSelectionCriteria.builder()
										.requireResidentKey(false)
										.build())
								.build(),
						new ByteArray(Base64.getDecoder()
								.decode(registrationRequestEntity.getChallenge()))));
		return request;
	}

	/**
	 * 开始注册
	 * 
	 * @param username
	 * @param displayName
	 * @param credentialNickname
	 * @param requireResidentKey
	 * @return
	 * @throws MalformedURLException
	 * @throws ExecutionException
	 * @throws JsonProcessingException
	 */
	public String startRegistration(
			@NonNull String username,
			Optional<String> displayName,
			Optional<String> credentialNickname,
			boolean requireResidentKey) throws MalformedURLException, ExecutionException,
			JsonProcessingException {
		logger.trace("startRegistration username: {}, credentialNickname: {}", username,
				credentialNickname);
		RegistrationRequestEntity registrationRequestEntity = new RegistrationRequestEntity(
				username, displayName.get(),
				credentialNickname.get(),
				Base64.getEncoder().encodeToString(UranusUtil.generateRandomByte(32)),
				Base64.getEncoder().encodeToString(UranusUtil.generateRandomByte(32)));
		RegistrationRequest request = this.getRegistrationRequest(registrationRequestEntity);
		// 存储到缓存 防止重放攻击
		cacheManager.getRequestCache().put(request.getRequestId().getBase64Url(),
				registrationRequestEntity);
		return this.writeJson(new StartRegistrationResponse(request));

	}

	private final class StartRegistrationResponse {

		public final boolean success = true;

		public final RegistrationRequest request;

		private StartRegistrationResponse(RegistrationRequest request)
				throws MalformedURLException {
			this.request = request;
		}
	}

	private String writeJson(Object o) throws JsonProcessingException {
		return jsonMapper.writeValueAsString(o);
	}

	@Value
	public static class SuccessfulRegistrationResult {

		final boolean success = true;

		RegistrationRequest request;

		RegistrationResponse response;

		CredentialRegistration registration;

		boolean attestationTrusted;

		Optional<AttestationCertInfo> attestationCert;

		@JsonSerialize(using = AuthDataSerializer.class)
		AuthenticatorData authData;

		String username;

		ByteArray sessionToken;

		public SuccessfulRegistrationResult(RegistrationRequest request,
				RegistrationResponse response,
				CredentialRegistration registration, boolean attestationTrusted,
				ByteArray sessionToken) {
			this.request = request;
			this.response = response;
			this.registration = registration;
			this.attestationTrusted = attestationTrusted;
			attestationCert = Optional.ofNullable(
					response.getCredential().getResponse().getAttestation()
							.getAttestationStatement().get("x5c"))
					.map(certs -> certs.get(0))
					.flatMap((JsonNode certDer) -> {
						try{
							return Optional.of(new ByteArray(certDer.binaryValue()));
						} catch (IOException e){
							logger.error("Failed to get binary value from x5c element: {}", certDer,
									e);
							return Optional.empty();
						}
					})
					.map(AttestationCertInfo::new);
			this.authData = response.getCredential().getResponse().getParsedAuthenticatorData();
			this.username = request.getUsername();
			this.sessionToken = sessionToken;
		}

	}

	@Value
	public static class AttestationCertInfo {

		final ByteArray der;

		final String text;

		public AttestationCertInfo(ByteArray certDer) {
			der = certDer;
			X509Certificate cert = null;
			try{
				cert = CertificateParser.parseDer(certDer.getBytes());
			} catch (CertificateException e){
				logger.error("Failed to parse attestation certificate");
			}
			if (cert == null){
				text = null;
			} else{
				text = cert.toString();
			}
		}
	}

	/**
	 * 完成注册
	 * 
	 * @author dell
	 * @param responseJson
	 *            认证器返回的结果集，包含公钥信息
	 */
	public boolean finishRegistration(String responseJson) {
		logger.info("finishRegistration responseJson: {}", responseJson);
		RegistrationResponse response = null;
		try{
			response = jsonMapper.readValue(responseJson, RegistrationResponse.class);
		} catch (IOException e){
			logger.error("JSON error in finishRegistration; responseJson: {}", responseJson, e);
			return false;
		}
		RegistrationRequestEntity registrationRequestEntity = cacheManager.getRequestCache().asMap()
				.get(response.getRequestId().getBase64Url());
		RegistrationRequest request = null;
		if (registrationRequestEntity != null){
			request = this.getRegistrationRequest(registrationRequestEntity);
		}

		if (request == null){
			logger.debug("fail finishRegistration responseJson: {}", responseJson);
			return false;
		} else{
			try{
				// 核心校验步骤
				RegistrationResult registration = this.getRelyingParty(getUserStorage())
						.finishRegistration(
								FinishRegistrationOptions.builder()
										.request(request.getPublicKeyCredentialCreationOptions())
										.response(response.getCredential())
										.build());
				return true;
			} catch (RegistrationFailedException e){
				logger.debug("fail finishRegistration responseJson: {}", responseJson, e);
				return false;
			} catch (Exception e){
				logger.error("fail finishRegistration responseJson: {}", responseJson, e);
				return false;
			}
		}
	}

	private String val(Map<String, Object> map, String key) {
		if (map == null || !map.containsKey(key)){
			return "";
		}
		return String.valueOf(map.get(key));
	}

	/**
	 * 开始认证
	 * 
	 * @param username
	 * @return
	 */
	public String startAuthentication(String username) {
		logger.trace("startAuthentication username: {}", username);
		if (StringUtils.isNotBlank(username)){
			RegistrationRequestEntity registrationRequestEntity = new RegistrationRequestEntity(
					username, username,
					username, Base64.getEncoder().encodeToString(UranusUtil.generateRandomByte(32)),
					Base64.getEncoder().encodeToString(UranusUtil.generateRandomByte(32)));
			List<PublicKeyCredentialDescriptor> allowCredentials = Lists.newArrayList();
			List<String> fidoRegisterInfos = cacheManager.getRegisterCache().asMap().get(username);
			for (String fidoRegisterInfo : fidoRegisterInfos){
				RegistrationResponse response = null;
				try{
					response = jsonMapper.readValue(fidoRegisterInfo, RegistrationResponse.class);
					allowCredentials.add(PublicKeyCredentialDescriptor.builder()
							.id(response.getCredential().getResponse().getAttestation()
									.getAuthenticatorData()
									.getAttestedCredentialData().get().getCredentialId())
							.build());
				} catch (IOException e){
					logger.error("JSON error in finishRegistration; responseJson: {}",
							fidoRegisterInfo, e);
				}
			}
			if (allowCredentials.isEmpty()){
				return null;
			}
			// 构造前端认证器需要的格式参数
			AssertionRequestWrapper request = new AssertionRequestWrapper(
					new ByteArray(Base64.getDecoder()
							.decode(registrationRequestEntity.getUserIdentityId())),
					this.getRelyingParty(getUserStorage()).startAssertion(
							StartAssertionOptions.builder()
									.username(username)
									.build(),
							new ByteArray(Base64.getDecoder()
									.decode(registrationRequestEntity.getChallenge())),
							allowCredentials));
			// 存储到缓存 防止重放攻击
			cacheManager.getRequestCache().put(request.getRequestId().getBase64Url(),
					registrationRequestEntity);
			try{
				return this.writeJson(new StartAuthenticationResponse(request));
			} catch (Exception e){
				logger.error("start assertion error : ", e);
			}
		}
		return null;
	}

	private final class StartAuthenticationResponse {

		public final boolean success = true;

		public final AssertionRequestWrapper request;

		private StartAuthenticationResponse(AssertionRequestWrapper request)
				throws MalformedURLException {
			this.request = request;
		}
	}

	@Value
	@AllArgsConstructor
	public static final class SuccessfulAuthenticationResult {

		private final boolean success = true;

		private final AssertionRequestWrapper request;

		private final AssertionResponse response;

		private final Collection<CredentialRegistration> registrations;

		@JsonSerialize(using = AuthDataSerializer.class)
		AuthenticatorData authData;

		private final String username;

		private final List<String> warnings;

		public SuccessfulAuthenticationResult(AssertionRequestWrapper request,
				AssertionResponse response,
				Collection<CredentialRegistration> registrations, String username,
				List<String> warnings) {
			this(
					request,
					response,
					registrations,
					response.getCredential().getResponse().getParsedAuthenticatorData(),
					username,
					warnings);
		}
	}

	/**
	 * 完成认证
	 * 
	 * @param responseJson
	 * @param username
	 * @return
	 */
	public boolean finishAuthentication(String responseJson, String username) {
		logger.info("finishAuthentication responseJson: {}", responseJson);

		final AssertionResponse response;
		try{
			response = jsonMapper.readValue(responseJson, AssertionResponse.class);
		} catch (IOException e){
			logger.error("Failed to decode response object", e);
			return false;
		}
		RegistrationRequestEntity registrationRequestEntity = cacheManager.getRequestCache().asMap()
				.get(response.getRequestId().getBase64Url());
		AssertionRequestWrapper request = null;
		List<String> fidoRegisterInfos = cacheManager.getRegisterCache().asMap().get(username);
		if (registrationRequestEntity != null){
			List<PublicKeyCredentialDescriptor> allowCredentials = Lists.newArrayList();
			// 构造allowCredentials，前端设备会自动识别对应的id
			for (String fidoRegisterInfo : fidoRegisterInfos){
				RegistrationResponse responseLdap = null;
				try{
					responseLdap = jsonMapper.readValue(fidoRegisterInfo,
							RegistrationResponse.class);
					allowCredentials.add(PublicKeyCredentialDescriptor.builder()
							.id(responseLdap.getCredential().getResponse().getAttestation()
									.getAuthenticatorData()
									.getAttestedCredentialData().get().getCredentialId())
							.build());
				} catch (IOException e){
					logger.error("JSON error in finishRegistration; responseJson: {}",
							fidoRegisterInfo, e);
				}
			}
			if (allowCredentials.isEmpty()){
				return false;
			}
			// 构造认证过程需要的数据格式
			request = new AssertionRequestWrapper(
					new ByteArray(Base64.getDecoder()
							.decode(registrationRequestEntity.getUserIdentityId())),
					this.getRelyingParty(getUserStorage()).startAssertion(
							StartAssertionOptions.builder()
									.username(username)
									.build(),
							new ByteArray(Base64.getDecoder()
									.decode(registrationRequestEntity.getChallenge())),
							allowCredentials));
		}
		boolean flag = false;
		if (request != null){
			// 绑定多个设备，遍历认证，
			for (String fidoRegisterInfo : fidoRegisterInfos){
				try{
					ByteArray userHandle = UranusUtil.generateRandom(32);
					if (response.getCredential().getResponse().getUserHandle().isPresent()){
						userHandle = response.getCredential().getResponse().getUserHandle().get();
					}
					AssertionResult result = this
							.getRelyingParty(
									getAssertUserStorage(username,
											fidoRegisterInfo, userHandle))
							.finishAssertion(
									FinishAssertionOptions.builder()
											.request(request.getRequest())
											.response(response.getCredential())
											.build());
					if (result.isSuccess()){
						flag = true;
					}
					if (flag){
						break;
					}
				} catch (Exception e){
					logger.error("Assertion failed", e);
				}
			}
		}
		return flag;
	}

	@Value
	public static final class DeregisterCredentialResult {

		boolean success = true;

		CredentialRegistration droppedRegistration;

		boolean accountDeleted;
	}

	static ByteArray rawEcdaKeyToCose(ByteArray key) {
		final byte[] keyBytes = key.getBytes();

		if (!(keyBytes.length == 64 || (keyBytes.length == 65 && keyBytes[0] == 0x04))){
			throw new IllegalArgumentException(String.format(
					"Raw key must be 64 bytes long or be 65 bytes long and start with 0x04, was %d bytes starting with %02x",
					keyBytes.length,
					keyBytes[0]));
		}

		final int start = keyBytes.length == 64 ? 0 : 1;

		Map<Long, Object> coseKey = new HashMap<>();

		coseKey.put(1L, 2L); // Key type: EC
		coseKey.put(3L, COSEAlgorithmIdentifier.ES256.getId());
		coseKey.put(-1L, 1L); // Curve: P-256
		coseKey.put(-2L, Arrays.copyOfRange(keyBytes, start, start + 32)); // x
		coseKey.put(-3L, Arrays.copyOfRange(keyBytes, start + 32, start + 64)); // y

		return new ByteArray(CBORObject.FromObject(coseKey).EncodeToBytes());
	}

	private static class AuthDataSerializer extends JsonSerializer<AuthenticatorData> {

		@Override
		public void serialize(AuthenticatorData value, JsonGenerator gen,
				SerializerProvider serializers)
				throws IOException {
			gen.writeStartObject();
			gen.writeStringField("rpIdHash", value.getRpIdHash().getHex());
			gen.writeObjectField("flags", value.getFlags());
			gen.writeNumberField("signatureCounter", value.getSignatureCounter());
			value.getAttestedCredentialData().ifPresent(acd -> {
				try{
					gen.writeObjectFieldStart("attestedCredentialData");
					gen.writeStringField("aaguid", acd.getAaguid().getHex());
					gen.writeStringField("credentialId", acd.getCredentialId().getHex());
					gen.writeStringField("publicKey", acd.getCredentialPublicKey().getHex());
					gen.writeEndObject();
				} catch (IOException e){
					throw new RuntimeException(e);
				}
			});
			gen.writeObjectField("extensions", value.getExtensions());
			gen.writeEndObject();
		}
	}
}
